import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


dataset = torchvision.datasets.CIFAR10(root='./data', train=False, transform=torchvision.transforms.ToTensor(), download=True)

dataloader = DataLoader(dataset, batch_size=1)

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.relu = nn.ReLU(False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input):
        output = self.relu(input)
        return output

tudui = Tudui()

writer = SummaryWriter('logs')

step = 0
for data in dataloader:
    imgs, targets = data

    writer.add_images("input", imgs, step)
    print(imgs.data)

    output = tudui(imgs)
    writer.add_images("output", output, step)

    step += 1

writer.close()
